#pragma once

template<typename T>
class CArray2D
{
public:
	CArray2D(UINT uiWidth, UINT uiHeight)
		: m_uiWidth(uiWidth)
		, m_uiHeight(uiHeight)
	{
		m_uiPitch = (sizeof(T) * uiWidth + 15) & (~15);
		m_pBuffer = (T*)_aligned_malloc(m_uiPitch * m_uiHeight, 16);
		m_bAllocateMemory = TRUE;
	}

	CArray2D(T* pBuffer, UINT uiWidth, UINT uiHeight, UINT uiPitch = sizeof(T) * uiWidth)
		: m_pBuffer(pBuffer)
		, m_uiWidth(uiWidth)
		, m_uiHeight(uiHeight)
		, m_uiPitch(uiPitch)
		, m_bAllocateMemory(FALSE)
	{
	}

	~CArray2D()
	{
		if(m_bAllocateMemory)
		{
			_aligned_free(m_pBuffer);
		}
	}

	T* operator [](UINT y)
	{
		return (T*)(((LPBYTE)m_pBuffer) + y * m_uiPitch);
	}

	const T* operator [](UINT y) const
	{
		return (const T*)(((LPBYTE)m_pBuffer) + y * m_uiPitch);
	}

	const T& Get(UINT y, UINT x) const
	{
		return ((T*)(((LPBYTE)m_pBuffer) + y * m_uiPitch))[x];
	}

	void Set(UINT y, UINT x, const T& t)
	{
		((T*)(((LPBYTE)m_pBuffer) + y * m_uiPitch))[x] = t;
	}

	void Clear(const T& t)
	{
		for(UINT i = 0; i < m_uiHeight; ++i)
		{
			for(UINT j = 0; j < m_uiWidth; ++i)
			{
				Set(i, j, t);
			}
		}
	}

	UINT GetWidth() const { return m_uiWidth; }
	UINT GetHeight() const { return m_uiHeight; }

private:
	T* m_pBuffer;
	UINT m_uiWidth;
	UINT m_uiHeight;
	UINT m_uiPitch;
	BOOL m_bAllocateMemory;
};

template<typename T>
class CArray3D
{
public:
	CArray3D(UINT uiWidth, UINT uiHeight, UINT uiDepth)
		: m_uiWidth(uiWidth)
		, m_uiHeight(uiHeight)
		, m_uiDepth(uiDepth)
	{
		m_uiPitch = (sizeof(T) * uiWidth + 15) & (~15);
		m_uiSlice = m_uiPitch * uiHeight;
		m_pBuffer = (T*)_aligned_malloc(m_uiSlice * m_uiDepth, 16);
		m_bAllocateMemory = TRUE;

		m_pSubArray = (CArray2D<T>**)new void*[m_uiDepth];
		for(UINT i = 0; i < m_uiDepth; ++i)
		{
			m_pSubArray[i] = new CArray2D<T>(
				(T*)(((LPBYTE)m_pBuffer) + i * m_uiSlice), 
				m_uiWidth, 
				m_uiHeight, 
				m_uiPitch);
		}
	}

	CArray3D(T* pBuffer, UINT uiWidth, UINT uiHeight, UINT uiDepth, UINT uiPitch = sizeof(T) * uiWidth, UINT uiSlice = sizeof(T) * uiWidth * uiHeight)
		: m_pBuffer(pBuffer)
		, m_uiWidth(uiWidth)
		, m_uiHeight(uiHeight)
		, m_uiDepth(uiDepth)
		, m_uiPitch(uiPitch)
		, m_uiSlice(uiSlice)
		, m_bAllocateMemory(FALSE)
	{
		m_pSubArray = (CArray2D<T>**)new void*[m_uiDepth];;
		for(UINT i = 0; i < m_uiDepth; ++i)
		{
			m_pSubArray[i] = new CArray2D<T>(
				(T*)(((LPBYTE)m_pBuffer) + i * m_uiSlice), 
				m_uiWidth, 
				m_uiHeight, 
				m_uiPitch);
		}
	}

	~CArray3D()
	{
		for(UINT i = 0; i < m_uiDepth; ++i)
		{
			delete m_pSubArray[i];
		}
		delete[] m_pSubArray;

		if(m_bAllocateMemory)
		{
			_aligned_free(m_pBuffer);
		}
	}

	CArray2D<T>& operator [](UINT z)
	{
		return *m_pSubArray[z];
	}

	const CArray2D<T>& operator [](UINT z) const
	{
		return *m_pSubArray[z];
	}

	const T& Get(UINT z, UINT y, UINT x) const
	{
		return m_pSubArray[z]->Get(y, x);
	}

	void Set(UINT z, UINT y, UINT x, const T& t)
	{
		return m_pSubArray[z]->Set(y, x, t);
	}

	void Clear(const T& t)
	{
		for(UINT i = 0; i < m_uiDepth; ++i)
		{
			for(UINT j = 0; j < m_uiHeight; ++j)
			{
				for(UINT k = 0; k < m_uiWidth; ++k)
				{
					Set(i, j, k, t);
				}
			}
		}
	}

	UINT GetWidth() const { return m_uiWidth; }
	UINT GetHeight() const { return m_uiHeight; }
	UINT GetDepth() const { return m_uiDepth; }

private:
	T* m_pBuffer;
	UINT m_uiWidth;
	UINT m_uiHeight;
	UINT m_uiDepth;
	UINT m_uiPitch;
	UINT m_uiSlice;
	CArray2D<T>** m_pSubArray;
	BOOL m_bAllocateMemory;
};